﻿#define VERBOSE
#define RandomWavelength
using System;
using System.Collections.Generic;
using System.Linq;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.Library.Data.Imaging;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.PathTracer;
using RayDen.RayEngine.Samplers;

namespace RayDen.RayEngine.Engines.XP
{
    public class SingleWavelengthPathSampler : PathSamplerBase
    {
        internal RgbSpectrum Throughput;
        public PathTracerPathState PathState;

        protected int depth, tracedShadowRayCount;
        protected float pathWeight, HeroWavelength, HWWeight, waveThroughput, waveRadiance, rs;
        protected bool specularBounce;

        public UInt64 PathExpression;


        protected RayEngineScene scene;
        protected internal SpectralSamplingHelper wlSampler;
        private ShadowRayInfo[] secRays;
        SurfaceIntersectionData hitInfo = null;
        public int MaxRaysPerPath
        {
            get;
            set;
        }


        private float[] lambdas, wls;
        private int lambdaSample;

        public SingleWavelengthPathSampler()
        {
            lambdas = new float[SpectralSamplingHelper.SpectralSamples];
            wls = new float[SpectralSamplingHelper.SpectralSamples];
            rs = 0;
        }



        public override void Splat(SampleBuffer sampleBuffer)
        {

            var waveLengthWeight = 1f/ (SpectralSamplingHelper.SpectralSamples);
            lambdas[lambdaSample] = (HeroWavelength);
            wls[lambdaSample] = (waveRadiance / waveLengthWeight);
            rs += waveRadiance;
            lambdaSample++;
            if (lambdaSample >= SpectralSamplingHelper.SpectralSamples)
            {
                RgbSpectrum pix = RgbSpectrum.ZeroSpectrum();
                if (rs > 0f)
                {
#if RandomWavelength

                var lb = lambdas.ToList();
                lb.Sort();
                lambdas = lb.ToArray();

#endif
                    var spd = new IrregularSPD(lambdas, wls, SpectralSamplingHelper.SpectralSamples, SpectralSamplingHelper.SpectralResolution, SPDResamplingMethod.Linear);      
                    //var spd = new RegularSPD(wls, SampledSpectrum.sampledLambdaStart, SampledSpectrum.sampledLambdaEnd,SpectralSamplingHelper.SpectralSamples);
                    spd.Normalize();
                    pix =
                        //ColorFactory.ToRgb(spd);
                        spd.ToRgb();
                }
                sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref pix);
                rs = 0;
                lambdaSample = 0;
            }
            this.InitPath(pathIntegrator);
        }


        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;
            this.Radiance = new RgbSpectrum(0f);
            this.waveRadiance = 0f;
            this.waveThroughput = 1f;
            this.Throughput = new RgbSpectrum(1f);
            this.PathState = PathTracerPathState.EyeVertex;
            if (this.secRays == null)
                this.secRays = new ShadowRayInfo[scene.ShadowRaysPerSample];

            if (lambdaSample == 0)
            {
                this.Sample = pathIntegrator.Sampler.GetSample(this.Sample);
                this.Sample.InitSample(scene.MaxPathDepth * 8);
                pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out this.PathRay);
            }

            this.HeroWavelength = this.wlSampler.SampleWavelength(
#if RandomWavelength

                Sample.GetLazyValue());
#else
lambdaSample);

#endif
            this.RayIndex = -1;
            this.pathWeight = 1.0f;
            this.tracedShadowRayCount = 0;
            this.depth = 0;
            this.specularBounce = true;
        }

        public override sealed void Advance(RayBuffer rayBuffer, SampleBuffer consumer)
        {
            int currentTriangleIndex = 0;
#if VERBOSE

            try
            {
#endif

                if (((PathState == PathTracerPathState.ShadowRaysOnly) || (PathState == PathTracerPathState.NextVertex)) &&
                    (tracedShadowRayCount > 0))
                {
                    var continueRays = new List<ShadowRayInfo>();
                    for (int i = 0; i < tracedShadowRayCount; ++i)
                    {
                        RayHit shadowRayHit = rayBuffer.rayHits[secRays[i].ShadowRayIndex];
                        RgbSpectrum attenuation;
                        bool continueTrace;
                        if (this.ShadowRayTest(ref shadowRayHit, ref secRays[i].ShadowRay, out attenuation, out continueTrace))
                        {
                            //Radiance += attenuation * ((secRays[i].Throughput) / secRays[i].Pdf);
                            waveRadiance += secRays[i].Throughput[0] / secRays[i].Pdf;
                            //pathWeight *= secRays[i].Pdf;
                        }

                        if (continueTrace)
                        {
                            continueRays.Add(secRays[i]);
                        }
                    }

                    if (PathState == PathTracerPathState.ShadowRaysOnly)
                    {
                        Splat(consumer);
                        return;
                    }

                    Array.Copy(continueRays.ToArray(), secRays, continueRays.Count);
                    tracedShadowRayCount = continueRays.Count;
                }
                RayHit rayHit = rayBuffer.rayHits[RayIndex];
                Vector wo = -PathRay.Dir;

                depth++;
                bool missed = rayHit.Index == 0xffffffffu;
                if (missed || PathState == PathTracerPathState.ShadowRaysOnly || depth > scene.MaxPathDepth)
                {
                    if (missed)
                    {
                        //Radiance += this.scene.SampleEnvironment(ref wo) * Throughput;
                        //var sampledEnvironment = this.scene.SampleEnvironment(ref wo); //!
                        waveRadiance += wlSampler.SampleEnvironment(HeroWavelength, ref wo);
                        //Radiance.MAdd(ref sampledEnvironment, ref Throughput);
                    }
                    Splat(consumer);

                    return;
                }

                //var mesh = scene.GetMeshByTriangleIndex(currentTriangleIndex);
                //rayHit.Index += (uint)mesh.StartTriangle;
                // Something was hit
                if (hitInfo == null)
                {
                    hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit);
                }
                else
                {
                    SurfaceSampler.GetIntersection(ref PathRay, ref rayHit, ref hitInfo);
                }

                currentTriangleIndex = (int)rayHit.Index;

                if (hitInfo == null)
                {
                    System.Diagnostics.Debugger.Break();
                }
                //If Hit light)

                if (hitInfo.IsLight)
                {
                    if (specularBounce || depth == 1)
                    {
                        var lt = scene.GetLightByIndex(currentTriangleIndex);

                        if (lt != null)
                        {
                            waveRadiance += wlSampler.SampleLight(lt, HeroWavelength) * waveThroughput;
                            //var le = hitInfo.Color * (RgbSpectrum)(lt.Le(ref wo)); //!
                            //Radiance += Throughput * le;
                            //Radiance.MAdd(ref Throughput, ref le);
                        }
                    }
                    Splat(consumer);

                    return;
                }

                var hitPoint = PathRay.Point(rayHit.Distance);

                tracedShadowRayCount = 0;
                var bsdf = wlSampler.GetBrdf(scene.GetMeshByTriangleIndex(currentTriangleIndex).MaterialID);


                //if (!hitInfo.TextureData.Alpha.IsBlack())
                //{
                //    Throughput *= (RgbSpectrum.UnitSpectrum() - hitInfo.TextureData.Alpha);
                //    PathRay = new RayData(hitPoint, -wo);
                //    return;
                //}



                if (bsdf.IsDiffuse())
                {
                    var kd = hitInfo.Color;
                    bsdf.Kd = ColorFactory.CreateRegularSpd(ref kd, SpectrumType.Reflectance);

                    float lightStrategyPdf = LightSampler.StrategyPdf;
                    float ltThroughput = waveThroughput;//* bsdf.Kd.Sample(HeroWavelength);
                    LightSampler.EvaluateShadow(ref hitPoint, ref hitInfo.Normal, Sample.GetLazyValue(depth),
                                                Sample.GetLazyValue(2 * depth + 1), Sample.GetLazyValue(3 * depth + 2), ref ls);
                    for (int index = 0; index < ls.Length; index++)
                    {

                        if (ls[index].Pdf <= 0f)
                            continue;
                        ls[index].Lambda = wlSampler.SampleLight(scene.Lights[ls[index].LightIndex], HeroWavelength);
                        secRays[tracedShadowRayCount].Throughput[0] = ls[index].Lambda;
                        secRays[tracedShadowRayCount].Pdf = ls[index].Pdf;
                        secRays[tracedShadowRayCount].ShadowRay = ls[index].LightRay;
                        Vector lwi = secRays[tracedShadowRayCount].ShadowRay.Dir;
                        float fs;
                        bsdf.F(ref secRays[tracedShadowRayCount].ShadowRay.Dir, ref wo, ref hitInfo.ShadingNormal, HeroWavelength, out fs);
                        secRays[tracedShadowRayCount].Throughput[0] *= ltThroughput *
                                                                    Vector.AbsDot(ref hitInfo.Normal, ref lwi) *
                                                                    fs;
                        if (secRays[tracedShadowRayCount].Throughput[0] > 0f)
                        {
#if DEBUG
                                    RayDen.Library.Core.Components.Assert.IsNotNaN(secRays[tracedShadowRayCount].Pdf);
#endif
                            secRays[tracedShadowRayCount].Pdf /= lightStrategyPdf;
                            tracedShadowRayCount++;
                        }
                    }
                }

                float fPdf = 0f;
                var wi = new Vector();
                float f;
                BsdfSampleData bsdfSample;
                hitInfo.TextureData.Throughput = Throughput;
                bsdf.Sample_f(ref wo, ref hitInfo.Normal, ref hitInfo.ShadingNormal, HeroWavelength, ref hitInfo.TextureData,
                                    Sample.GetLazyValue(4 * depth + 3), Sample.GetLazyValue(5 * depth + 4), Sample.GetLazyValue(6 * depth + 5)
                                    , out bsdfSample);
                var color = ColorFactory.CreateRegularSpd(ref hitInfo.Color, SpectrumType.Reflectance);

                f = bsdfSample.Lambda;//*color.Sample(HeroWavelength)
                fPdf = bsdfSample.Pdf;
                wi = bsdfSample.Wi;
                specularBounce = bsdfSample.SpecularBounce;

                if ((fPdf <= 0.0f))// || f.IsBlack()
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        Splat(consumer);

                    }
                    return;
                }
                waveThroughput *= f / fPdf;
                pathWeight *= fPdf;
                //Throughput *= (f * hitInfo.Color) / fPdf;

                if (depth > scene.MaxPathDepth)
                {
                    float prob = Math.Max(waveThroughput, scene.RussianRuletteImportanceCap);
                    if (prob >= Sample.GetLazyValue(7 * depth + 5))
                    {
                        //Throughput /= prob;
                        waveThroughput /= prob;
                        pathWeight *= prob;
                    }
                    else
                    {
                        if (tracedShadowRayCount > 0)
                            PathState = PathTracerPathState.ShadowRaysOnly;
                        else
                        {
                            Splat(consumer);

                        }

                        return;
                    }
                }

                PathRay.Org = hitPoint;
                PathRay.Dir = wi.Normalize();
                PathState = PathTracerPathState.NextVertex;

#if VERBOSE
            }
            catch (Exception ex)
            {
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Advance path exception");
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Error triangle {0}", currentTriangleIndex);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.StackTrace);

            }
#endif
        }



        public override bool FillRayBuffer(RayBuffer rayBuffer)
        {
            var leftSpace = rayBuffer.LeftSpace();
            if (((PathState == PathTracerPathState.EyeVertex) && (1 > leftSpace)) ||
            ((PathState == PathTracerPathState.ShadowRaysOnly) && (tracedShadowRayCount > leftSpace)) ||
            ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount + 1 > leftSpace)))
                return false;
            if (PathState != PathTracerPathState.ShadowRaysOnly)
                RayIndex = rayBuffer.AddRay(ref PathRay);
            if (PathState == PathTracerPathState.NextVertex || PathState == PathTracerPathState.ShadowRaysOnly)
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    secRays[i].ShadowRayIndex = rayBuffer.AddRay(ref secRays[i].ShadowRay);
            }
            return true;
        }

        protected override bool ShadowRayTest(ref RayHit shadowRayHit, ref RayData shadowRay, out RgbSpectrum attenuation, out bool continueTrace)
        {
            var hit = shadowRayHit.Index == 0xffffffffu || scene.IsLight((int)shadowRayHit.Index);
            attenuation = new RgbSpectrum(1f);
            continueTrace = false;

            if (hit) return true;
            var mesh = scene.GetMeshByTriangleIndex((int)shadowRayHit.Index);
            var mat =
                //SurfaceMaterials.CreateMaterial(scene.MaterialProvider.Get(mesh.MaterialName));
                scene.MatLib.GetSurfMat(mesh != null ? mesh.MaterialName : "default");

            if (mat.TransparentShadows)
            {
                continueTrace = true;
                return true;
            }

            if (mat.AlphaTexture != null)
            {
                continueTrace = true;
                this.SurfaceSampler.GetIntersection(ref shadowRay, ref shadowRayHit, ref hitInfo, IntersectionOptions.ResolveTextures);
                attenuation = (RgbSpectrum)hitInfo.TextureData.Alpha;
                hit = true;
                //mat.MaterialData.Kt.Filter() < Sample.GetLazyValue();
            }

            return hit;
        }
    }
}
